# -*- coding: utf-8 -*-
import sys
import numpy as np
import pandas as pd
import datetime
import importlib

from tools.config import CONFIG
from tools.functions import *
from tools.aqi_compute import aqi_compute_from_hours


def filt_between_date(df, left, right):
    for date_range_left, date_range_right in CONFIG().DATASET_FILT_PARAMS['to_delete_qb_ranges']:
        df = df[(df['qb_time'] > date_range_right) & (df['qb_time'] < date_range_left)]
    return df
def load_data(station: str):
    xy = os.path.join(CONFIG().MERGED_DATASET_PATH_ZSCORE, '%s_MERGE.csv'%station)
    df_xy = pd.read_csv(xy)
    df_xy = df_xy.dropna(how='any', axis=0).reset_index(drop=True)
    df_xy['模型运行日期'] = pd.to_datetime(df_xy['模型运行日期'])
    
    df_test_xy = df_xy[df_xy['模型运行日期'].isin(CONFIG().EXAM_DATES)]# todo
    df_train_xy = df_xy[(df_xy['模型运行日期'] < CONFIG().EXAM_DATES[0])]
    assert len(df_test_xy) == 72
    
    X_train = df_train_xy[[i for i in list(df_train_xy) if i in CONFIG().FEATURES]]
    y_train = df_train_xy[CONFIG().LABELS]

    df_test_xy.to_csv('exam_data/exam_%s.csv'%station)
    return X_train, y_train
# X_train, X_test, y_train, y_test = load_data('A')


# +
def fit_save_model(station):
    print('start station:', station)
    X_train, y_train = load_data(station)
    model = EXP_MODEL(station=station)
    performance = model.fit(X_train, y_train)
    return performance

def merge_results_to_csv(result_dict):
    ret = pd.DataFrame()
    for station in  result_dict:
        df = result_dict[station]
        df = df[[i for i in list(df) if i not in ['监测时间', '监测日期_pred']]]
        df['station'] = station
        ret = pd.concat([ret,df], axis=0)
    ret = ret[['station']+[i for i in list(ret) if (('Unnamed' not in i) and ('station'!=i))]]
    ret.to_csv('exam_data/aqi_score.csv')
    return ret
    
def predict(station):
    df_test_xy = pd.read_csv('exam_data/exam_%s.csv'%station)
    X_test = df_test_xy[[i for i in list(df_test_xy) if i in CONFIG().FEATURES]]
    y_test = df_test_xy[CONFIG().LABELS]
    model = EXP_MODEL(station=station)
    y_pred = model.predict(X_test)
    ret1 = pd.DataFrame()
    ret2 = None
    for y_label in CONFIG().LABELS:
        # 拟合效果图
        show_figure(y_test[y_label], y_pred[y_label], 
                    title=station+'-'+y_label+'test-fitness', 
                    save_path='exam_data/fitness_%s_%s.png'%(station, y_label))
        # MSE效果
        scores = {'场站':[station], 'label':y_label,}
        scores.update(my_score(y_pred[y_label], y_test[y_label], 'test_'))
        ret1 = pd.concat([ret1, pd.DataFrame(scores)])
    # AQI效果
    time_name = '监测时间' if '监测时间' in list(df_test_xy) else '实测时间'
    df_y_pred = pd.concat([df_test_xy[time_name], y_pred],axis=1)
    df_aqi_test = aqi_compute_from_hours(df_test_xy, station)
    df_aqi_pred = aqi_compute_from_hours(df_y_pred, station)
    ret2 = df_aqi_test.merge(df_aqi_pred, how='left', on=time_name, suffixes=('','_pred'))
#     ret2.to_csv('exam_data/aqi_score_%s.csv'%(station))
    return ret2


# -

if 1 and __name__ == "__main__":
    multi_result = pd.DataFrame()
    i = 0
    STATIONS = ['A', 'B', 'C']# CONFIG().STATIONS
    EXP_ROOTs = [os.path.join(CONFIG().MODEL_ROOT, i) for i in os.listdir(CONFIG().MODEL_ROOT) if 'model' in i]
    TRAIN = False#True
    for exp_root in EXP_ROOTs:
        print('[Model] ------> ',str(exp_root).split('/')[-1])
        module_name = 'models.' + str(exp_root).split('/')[-1] + '.' + str(exp_root).split('/')[-1]
        print(module_name)
        EXP_MODEL = getattr(importlib.import_module(module_name), 'PyModel')        
        if TRAIN:
            print('train start', datetime.datetime.now())
#             performances = functions.pool(
#                 fit_save_model, [(station,) for station in STATIONS]
#             )
            performances = [fit_save_model(station) for station in STATIONS]
            performances = pd.concat(performances, axis=0)
            performances.to_csv(os.path.join(exp_root, 'performances.csv'))
    # 计算测试集结果
    print('load and predict start', datetime.datetime.now())
#     pred_results_ = pool(predict, [(station,) for station in STATIONS])
    pred_results_ = {station:predict(station,) for station in STATIONS}
    merge_results_to_csv(pred_results_)
    print('all ok.')
